
import os
# 选用RolePlay 配置agent
from modelscope_agent.agents.role_play import RolePlay 
from tools import *
from modelscope.utils.config import Config
import json
from raw_prompt import ROLE_TEMPLATE
import re
from pprint import pprint

# 配置环境变量
os.environ['DASHSCOPE_API_KEY']="sk-548330d34bbf4650875532ce766e978e"
# os.environ['AMAP_TOKEN']="b84f785a1aab09acb0e620b5fa94eda6"
# tool_cfg_file = './cfg/tool_config.json'

class IndustryAI:
    '''工业AI类'''

    def __init__(self) -> None:
        ## 提示词模板
        self.role_template = ROLE_TEMPLATE
        
        ## 大模型配置
        self.llm_config = {  
                    'model'       :  'qwen-max', # 千问
                    'model_server':  'dashscope'
                    }
    
        ## 工具列表
        self.function_list = []

        self.bot = RolePlay(
                    function_list = self.function_list, 
                            llm = self.llm_config, 
                    instruction = self.role_template, )
        self.res = {
                    'model_res': '',
                    'dict_res' : { 'task': '', 'paras': {},},
                    'warn' : ''}
        self.task_paras_num = {'deviceInfo'       : 1,  # 记录每种task的参数个数
                               'workorderSchedule': 1,
                               'postponeReason'   : 1,
                               'materialStoreInfo': 1,
                               'insertWorkOrder'  : 4,
                               }
        
    
    def parse(self, question:str='')-> str:

        try:
            # 机器人运行agent
            response = self.bot.run(question)

            model_res = ''
            for chunk in response:
                model_res += chunk
            self.res['model_res'] = model_res

            if 'Error' in model_res:
                self.res['warn'] = '访问模型发生报错!'
                self.res['dict_res'] = {}
                
                return self.res
            if '无法归类' in model_res:
                self.res['warn'] = '该问题无法归类!'
                self.res['dict_res'] = {}

                return self.res

            dict_res = extract_to_dict(model_res)
            if not dict_res['task']:
                self.res['warn'] = '未知原因,大模型未识别出任务种类!'
                self.res['dict_res'] = {}
                return self.res
            
            warn = self.check_paras_num(dict_res) # 检验参数的个数

            self.res = {
                'dict_res' : dict_res,
                'model_res': model_res,
                'warn' : warn
            }

            # pprint("res:",self.res)
            return self.res
        
        except Exception as e:
            print(f'OMG! 发生了错误:{e}')
            self.res['warn'] = f'Agent内部异常: {str(e)[:50]}...(略)'
            return self.res
        
    def check_paras_num(self,dict_res)-> str:
        tpn = self.task_paras_num
        task = dict_res['task']
        num = tpn[task]
        provide = len(dict_res['paras'])
        return f"任务{task}需要{num}个参数,仅仅提供了{provide}个;" if num > provide else ''




def extract_to_dict(text:str='')-> dict:
    '''把大模型的结果中,重要的关键词/参数解析成字典格式'''

    dict_res= { 'task': '',
                'paras': {},
                }
    # 使用正则表达式匹配 [ ] 中的内容
    pattern = r'\[([^\]]+)\]'
    matches = re.findall(pattern, text)
    # print('matches:',matches)

    if not matches:
        return dict_res

    for match in matches:  # 将匹配到的字符串分割成键值对，并转换为字典
        key, val = match.split(':') # 分割键和值
        key, val = key.strip(), val.strip()
        if key in ['task','任务名称']:
            dict_res['task'] = val
        else:
            if val:
                dict_res['paras'][key] = val

    return dict_res

            
if __name__ == "__main__":

    # ma = "该问题所属种类为 [ 任务名称 : insertWorkOrder ], 该任务包含如下几个参数: 参数1:[ materialName : AAA螺母 ]; 参数2:[ wbslName : project32 ]; 参数4:[ num : 121 ];"
    # print('ma:',ma)
    # res = extract_to_dict(ma)

    # exit()

    iai = IndustryAI()
    print('准备提问！')
    # question = "请帮我查下火焰切割设备3的信息吧"
    # question = "请帮我查下处理进度，工单的号20241015121"
    # question = "我需要对物料AAA螺母插单, 项目是project32, 交货日期是今年3月五号, 数量是121件"
    question = "我需要对物料AAA螺母插单, 项目是project32"
    # question = "请我需要对物料插单"
    # question = "今年美国大选获胜者是谁"
 
    res = iai.parse(question)
    print(res)
    # print(json.dumps(res, indent=3))

    pass